import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from torch import nn

"""
线性回归模型详细写法
"""
data = pd.read_csv('dataset/Income.csv')
# 数据预处理，创建数据张量
X = torch.from_numpy(data.Education.values.reshape(-1, 1).astype(np.float32))
Y = torch.from_numpy(data.Income.values.reshape(-1, 1).astype(np.float32))

# 初始化w和b
# 单变量线性回归,并跟踪运算
w = torch.randn(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
"""
模型公式 ：  w@x + b 
"""
#  学习速率
learning_rate = 0.0001
#  对全部数据训练5000次
for epoch in range(5000):
    for x, y in zip(X, Y):
        # 等价于模型,权重去乘以x值+b得到一个预测值
        y_pred = torch.matmul(x, w) + b
        # 实际值和预测值查的平方求均值，【均方误差】
        loss = (y - y_pred).pow(2).mean()
        if w.grad is not None:
            w.grad.data.zero_()
        if b.grad is not None:
            b.grad.data.zero_()
        # 反向传播计算梯度，计算一个梯度值
        loss.backward()
        with torch.no_grad():
            # 下降梯度算法
            w.data -= w.grad.data * learning_rate
            b.data -= b.grad.data * learning_rate

